import os
import random
import shutil
import seedir as sd
import labelme2yolo_point as point
from tqdm import tqdm


def divide_files_gather(file_path, test_frac=0.2, file_extension='.jpg'):
    """
    划分测试集和训练集
    :param file_path:
    :param test_frac:
    :param file_extension:
    :return:
    """
    random.seed(123)  # 随机数种子，便于复现
    # 使用列表推导式筛选出指定后缀的文件名
    img_paths = [f for f in os.listdir(file_path) if f.endswith(file_extension)]
    random.shuffle(img_paths)  # 随机打乱
    val_number = int(len(img_paths) * test_frac)  # 测试集文件个数
    train_files = img_paths[val_number:]  # 训练集文件名列表
    val_files = img_paths[:val_number]  # 测试集文件名列表
    print('数据集文件总数', len(img_paths))
    print('训练集文件个数', len(train_files))
    print('测试集文件个数', len(val_files))
    os.makedirs('images/train', exist_ok=True)
    for each in tqdm(train_files):
        shutil.move(os.path.join(file_path, each), 'images/train')
    os.makedirs('images/val', exist_ok=True)
    for each in tqdm(val_files):
        shutil.move(os.path.join(file_path, each), 'images/val')
    os.makedirs('labels/train', exist_ok=True)
    for each in tqdm(train_files):
        base_name = os.path.splitext(each)[0]
        json_path = base_name + '.json'
        shutil.move(os.path.join(file_path, json_path), 'labels/train')
    os.makedirs('labels/val', exist_ok=True)
    for each in tqdm(val_files):
        base_name = os.path.splitext(each)[0]
        json_path = base_name + '.json'
        shutil.move(os.path.join(file_path, json_path), 'labels/val')
    sd.seedir(file_path, style='emoji', depthlimit=2)


def convert_json_txt(file_path):
    '''
    对当前目录的进行循环处理转换，然后删除.json
    :param file_path:
    :return:
    '''
    # 列出当前目录下的所有目录
    subdirectories = [os.path.join(file_path + '/labels', d) for d in os.listdir(file_path + '/labels') if
                      os.path.isdir(os.path.join(file_path + '/labels', d))]
    for txt_path in subdirectories:
        print('当前处理的文件夹地址为: ', txt_path)
        for labelme_path in [os.path.join(txt_path, f) for f in os.listdir(txt_path) if
                             f.endswith('.json') and os.path.isfile(os.path.join(txt_path, f))]:
            try:
                point.process_single_json(labelme_path)
                os.remove(labelme_path)
                print(f"已删除文件: {labelme_path}")
            except Exception as e:
                print('******有误******', labelme_path, e)
        print('YOLO格式的txt标注文件已保存至 ', txt_path)


if __name__ == '__main__':
    file_path = '/Users/admin/pywordspace/yolov8_demo/datasets'
    # divide_files_gather(file_path)
    convert_json_txt(file_path)
